# coding: utf-8

import math
import numpy as np
import matplotlib.pyplot as plt

# plt.style.use(["science", "no-latex"])

#####################################################################
############################ 多项式拟合  #############################
#####################################################################

# 功能：
#   给定一组X和Y数据, 进行多项式拟合
# 参数:
#   xy   -- 坐标点数据, 示例: [(x1, y1), (x2, y2), ... ]
#   n    -- 拟合次数
# 返回值:
#   列表, 示例: [-0.2506347664975213, 64.65774115325407, -1326.4090233350287]
# 注: 从高阶 -> 低阶, 例如: [a, b, c]
def polyfit(xy, n=2):
    return list(np.polyfit([x for x, _ in xy], [y for _, y in xy], n))


#####################################################################
############################ 抛物线计算  #############################
#####################################################################

# 功能：
#   已知(x1,y1) (x2,y2) (x3,y3)，求一条抛物线的参数a,b,c
# 方法:
#   解方程组：y1=f(x1);y2=f(x2);y3=f(x3), 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 3个点坐标, 示例: [(x1,y1), (x2,y2), (x3,y3)]
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_three_points(xy):
    x1, y1 = xy[0]
    x2, y2 = xy[1]
    x3, y3 = xy[2]
    a = (
        -1
        * ((y2 - y3) * x1 - (x2 - x3) * y1 + x2 * y3 - x3 * y2)
        / ((x2 - x3) * (x1 - x2) * (x1 - x3) + 1.0e-12)
    )
    b = (
        (y2 - y3) * x1 ** 2 + x2 ** 2 * y3 - x3 ** 2 * y2 - (x2 ** 2 - x3 ** 2) * y1
    ) / ((x2 - x3) * (x1 - x2) * (x1 - x3) + 1.0e-12)
    c = (
        (x2 * y3 - x3 * y2) * x1 ** 2
        - (x2 ** 2 - x3 ** 2) * x1
        + (x2 ** 2 * x3 - x2 * x3 ** 2) * y1
    ) / ((x2 - x3) * (x1 - x2) * (x1 - x3) + 1.0e-12)
    return a, b, c


# 功能：
#   已知2个点以及第1个点的导数求一条抛物线方程
# 方法:
#   解方程组：y1=f(x1);y2=f(x2);f'(x1)=k, 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 2个点坐标, 示例: [(x1,y1), (x2,y2)]
#   k -- 第1个点的导数
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_two_points_and_diff_1(xy, k):
    x1, y1 = xy[0]
    x2, y2 = xy[1]
    a = (k * (x1 - x2) - (y1 - y2)) / ((x1 - x2) ** 2 + 1.0e-12)
    b = -1 * (k * (x1 ** 2 - x2 ** 2) - 2 * x1 * (y1 - y2)) / ((x1 - x2) ** 2 + 1.0e-12)
    c = (
        -1
        * (k * x1 * x2 ** 2 - (k * x2 + y2) * x1 ** 2 + (2 * x1 * x2 - x2 ** 2) * y1)
        / ((x1 - x2) ** 2 + 1.0e-12)
    )
    return a, b, c


# 功能：
#   已知2个点以及第2点的导数求一条抛物线方程
# 方法:
#   解方程组：y1=f(x1);y2=f(x2);f'(x2)=k, 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 2个点坐标, 示例: [(x1,y1), (x2,y2)]
#   k -- 第2个点的导数
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_two_points_and_diff_2(xy, k):
    # 列表反转后再调用
    # xy[::-1] 等价于 list(reversed(xy))
    return solve_parabolic_two_points_and_diff_1(xy[::-1], k)


#####################################################################
##################### 正反抛物线相切、相交判定  #######################
#####################################################################

# 功能:
#   正反抛物线的交点个数
# 方法:
#   Δ = b^2-4ac
#   (1) Δ < 0  0个交点
#   (2) Δ = 0  1个交点
#   (3) Δ > 0  2个交点
# 参数:
#   a1, b1, c1 -- 正抛物线的2次、1次、0次系数
#   a2, b2, c2 -- 反抛物线的2次、1次、0次系数
# 返回值:
#   交点个数(取值范围: 0、1、2)
def get_two_parabolics_intersections(a1, b1, c1, a2, b2, c2):
    delta = (b1 - b2) ** 2 - 4 * (a1 - a2) * (c1 - c2)
    if delta < 0:
        # 无交点
        return 0
    elif abs(delta) < 1e-6:
        # 仅1个交点(其实也是切点)
        return 1
    else:
        # 2个交点
        return 2


# 功能:
#   求正反抛物线的交点或切点坐标
# 参数:
#   a1, b1, c1 -- 正抛物线的2次、1次、0次系数
#   a2, b2, c2 -- 反抛物线的2次、1次、0次系数
#   k          -- 当存在2个交点时, 取第k个交点(默认k=2)
# 返回值:
#   (1) 如果无交点, 则抛出异常(用户需自行处理, try/except)
#   (2) 如果有1个交点, 也等价于切点(仅限于正反抛物线的情况), 用元组类型(x0, y0)表示该切点并返回
#   (3) 如果有2个交点, 则返回第2个交点
def solve_two_parabolics_intersection_point(a1, b1, c1, a2, b2, c2, k=2):
    delta = (b1 - b2) ** 2 - 4 * (a1 - a2) * (c1 - c2)
    if delta < 0:
        # 情况1: 无交点
        raise Exception("无交点!")

    if abs(delta) < 1e-6:
        # 情况2: 只有1个交点(其实也是切点)
        x0 = -0.5 * (b1 - b2) / (a1 - a2 + 1.0e-12)
        return x0, a2 * x0 ** 2 + b2 * x0 + c2
    else:
        # 情况3: 存在2个交点
        x1 = -0.5 * ((b1 - b2) + math.sqrt(delta)) / (a1 - a2 + 1.0e-12)
        x2 = -0.5 * ((b1 - b2) - math.sqrt(delta)) / (a1 - a2 + 1.0e-12)
        # 确保x1 < x2
        if x1 > x2:
            # 交换x1和x2的值
            x1, x2 = x2, x1

        if k == 1:
            # 返回第1个交点
            return x1, a2 * x1 ** 2 + b2 * x1 + c2
        else:
            # 返回第2个交点
            return x2, a2 * x2 ** 2 + b2 * x2 + c2


#####################################################################
######################### 正反抛物线计算  ############################
#####################################################################

# 功能:
#   给定一组数据xy,共3个点坐标, 计算2条相切抛物线的参数
#   注: 两条抛物线必须满足相切条件
# 参数:
#   a2, b2, c2 -- 右抛物线
#   xy_min     -- 左抛物线的最小极值点(元组类型表示x、y坐标)
# 返回值:
#   a1, b1, c1, x0, y0 -- 左抛物线(前3个), 切点坐标(后2个)
def solve_tangent_parabolic(a2, b2, c2, xy_min):
    xmin, ymin = xy_min
    # 求左抛物线
    L1 = b2 ** 2 - 4 * a2 * c2 + 4 * a2 * ymin
    L2 = a2 * xmin ** 2 + b2 * xmin + c2 - ymin
    a1 = -1 * L1 / (4 * L2 + 1.0e-12)
    b1 = L1 * xmin / (2 * L2 + 1.0e-12)
    c1 = ymin - L1 * xmin ** 2 / (4 * L2 + 1.0e-12)
    x0, y0 = solve_two_parabolics_intersection_point(a1, b1, c1, a2, b2, c2)
    return a1, b1, c1, x0, y0


# 功能:
#   计算正反抛物线上任意一点的风压值
# 参数:
#   b2, b1, b0 -- 左抛物线
#   a2, a1, a0 -- 右抛物线
#   q0         -- 切点风量
#   q          -- 用户输入的风量
# 返回值:
#  正反抛物线上任意一点的风压值
def f0(b2, b1, b0, a2, a1, a0, q0, q):
    if q - q0 < 0:
        return b2 * q ** 2 + b1 * q + b0
    else:
        return a2 * q ** 2 + a1 * q + a0


# 功能: 算正反抛物线上任意一点的风压值一阶导
# 参数:
#   b2, b1, b0 -- 左抛物线
#   a2, a1, a0 -- 右抛物线
#   q0         -- 切点风量
#   q          -- 用户输入的风量
# 返回值:
#  正反抛物线上任意一点的风压导数
def f1(b2, b1, b0, a2, a1, a0, q0, q):
    if q - q0 < 0:
        return 2 * b2 * q + b1
    else:
        return 2 * a2 * q + a1


#####################################################################
###################### 绘图相关函数(测试用)  #########################
#####################################################################

# latex文本格式
def latex_str(s):
    return "$" + s + "$"


# 绘制一个点
def draw_point(
    plt, x0, y0, xmin=0, ymin=0, style="ro", text_offset_x=0, text_offset_y=0
):
    plt.plot([x0, x0], [ymin, y0], "--k", linewidth=0.5)
    plt.plot([xmin, x0], [y0, y0], "--k", linewidth=0.5)
    plt.plot(x0, y0, style)
    plt.text(
        x0 + text_offset_x, y0 + text_offset_y, "(%.2f, %.2f)" % (x0, y0), fontsize=9
    )


# 绘制正反抛物线、切点(x0,y0)、工况点points
def plot_parabolic(a1, b1, c1, a2, b2, c2, **kwargs):

    # 抛物线的绘制范围
    x1_min = kwargs.get("x1_min", 0)
    x1_max = kwargs.get("x1_max", 100)
    x2_min = kwargs.get("x2_min", 30)
    x2_max = kwargs.get("x2_max", 200)
    y_min = kwargs.get("y_min", 0)
    # 额外要绘制的点坐标
    points = kwargs.get("points", [])
    X1 = np.linspace(x1_min, x1_max, num=100)
    X2 = np.linspace(x2_min, x2_max, num=100)
    f1 = lambda x: a1 * x ** 2 + b1 * x + c1
    f2 = lambda x: a2 * x ** 2 + b2 * x + c2
    plt.figure()
    plt.plot(
        X1,
        [f1(x) for x in X1],
        "-b",
        label=latex_str("h1=%.4fq^2+%.4fq+%.4f" % (a1, b1, c1)),
    )
    plt.plot(
        X2,
        [f2(x) for x in X2],
        "-r",
        label=latex_str("h2=%.4fq^2+%.4fq+%.4f" % (a2, b2, c2)),
    )
    # 相切点
    # draw_point(
    #     plt,
    #     x0,
    #     y0,
    #     xmin=xmin,
    #     ymin=ymin,
    #     style="ks",
    #     text_offset_x=-15,
    #     text_offset_y=60,
    # )
    # [测试用]临时绘制抛物线的极值点
    x = -0.5 * b1 / (a1 + 1.0e-12)
    points.append((x, f1(x)))
    x = -0.5 * b2 / (a2 + 1.0e-12)
    points.append((x, f2(x)))
    for x, y in points:
        # 工况点
        draw_point(
            plt, x, y, xmin=x1_min, ymin=y_min,
        )

    # plt.title("x")
    plt.xlabel(latex_str("Q(m^3/min)"))
    plt.ylabel(latex_str("H(Pa)"))
    plt.legend(loc="upper right", fontsize=8)
    plt.ylim(ymin=y_min)
    plt.xlim(xmin=x1_min, xmax=x2_max)

    # 设置刻度及文字高度
    plt.tick_params(axis="both", which="major", labelsize=8)

    # 设置xy轴的刻度间距
    # 把x轴的刻度间隔设置为10，并存在变量里
    x_major_locator = plt.MultipleLocator(10)
    # 把y轴的刻度间隔设置为200，并存在变量里
    y_major_locator = plt.MultipleLocator(100)
    ax = plt.gca()
    ax.xaxis.set_major_locator(x_major_locator)
    ax.yaxis.set_major_locator(y_major_locator)

    # plt.grid(True)
    plt.show()


# 测试函数: 正反相切抛物线计算及绘制
def test_tangent_parabolics():
    # xy = [(110, 400.38), (151.19, 324.7)]
    # # # 已知2点坐标, 且第1点为极值点(导数为0), 计算抛物线
    # a, b, c = solve_parabolic_two_points_and_diff_1(xy, 0)
    # print("测试1: 3个点求抛物线 -> ", a, b, c)

    # print("\n测试2: 根据工况点、最大工况点、最小工况点求相切抛物线 -> ")
    # 工况点(右抛物线)、极大工况点(右抛物线)、极小工况点(左抛物线)
    # 求右抛物线(已知2点)
    # xy = [(151.19, 324.70), (110.00, 400.38)]
    # a2, b2, c2 = solve_parabolic_two_points_and_diff_2(xy, 0)

    # 已知右抛物线参数
    a2, a1, a0 = -0.04460643748459064, 9.813416246609941, -139.3578935635479

    # 求左抛物线和切点(已知最小工况点和右抛物线, 求相切抛物线)
    qh_min = (45.00, 228.217)
    b2, b1, b0, q0, h0 = solve_tangent_parabolic(a2, a1, a0, qh_min)
    print("左抛物线:", b2, b1, b0)
    print("右抛物线:", a2, a1, a0)
    print("切点:", q0, h0)

    # 测试: 任意给定风量计算压力
    points = [(q0, h0)]
    q = -10
    h = f0(b2, b1, b0, a2, a1, a0, q0, q)
    print("q=%.2f  h=%.2f" % (q, h))
    points.append((q, h))

    q = 10
    h = f0(b2, b1, b0, a2, a1, a0, q0, q)
    print("q=%.2f  h=%.2f" % (q, h))
    points.append((q, h))

    q = 70
    h = f0(b2, b1, b0, a2, a1, a0, q0, q)
    print("q=%.2f  h=%.2f" % (q, h))
    points.append((q, h))

    q = 120
    h = f0(b2, b1, b0, a2, a1, a0, q0, q)
    print("q=%.2f  h=%.2f" % (q, h))
    points.append((q, h))

    # 绘制抛物线
    plot_parabolic(b2, b1, b0, a2, a1, a0, points=points)


# 两条正反抛物线切的情况
def test1():
    # 正抛物线
    b2, b1, b0 = 0.47116293213937355, -42.40466389254428, 1182.3219375822314
    # 反抛物线
    a2, a1, a0 = -0.04460643748459064, 9.813416246609941, -139.3578935635479

    try:
        x0, y0 = solve_two_parabolics_intersection_point(b2, b1, b0, a2, a1, a0)
        print("交点或切点:", x0, y0)
        # 绘制抛物线
        plot_parabolic(b2, b1, b0, a2, a1, a0, points=[(x0, y0)], x1_min=0, x1_max=70)
    except Exception as e:
        print(e)


# 两条正反抛物线不相交的情况
def test2():
    # 正抛物线
    b2, b1, b0 = 0.57116293213937355, -42.40466389254428, 1182.3219375822314
    # 反抛物线
    a2, a1, a0 = -0.04460643748459064, 9.813416246609941, -139.3578935635479

    try:
        x0, y0 = solve_two_parabolics_intersection_point(b2, b1, b0, a2, a1, a0)
        print("交点或切点:", x0, y0)
        # 绘制抛物线
        plot_parabolic(b2, b1, b0, a2, a1, a0, points=[(x0, y0)])
    except Exception as e:
        # 打印异常信息
        print(e)


# 两条正反抛物线相交的情况
def test3():
    # 正抛物线
    b2, b1, b0 = 0.37116293213937355, -42.40466389254428, 1182.3219375822314
    # 反抛物线
    a2, a1, a0 = -0.04460643748459064, 9.813416246609941, -139.3578935635479

    try:
        # k=1 (表示取第1个交点)
        x1, y1 = solve_two_parabolics_intersection_point(b2, b1, b0, a2, a1, a0, k=1)
        print("第1个交点:", x1, y1)
        # 绘制抛物线
        plot_parabolic(b2, b1, b0, a2, a1, a0, points=[(x1, y1)], x1_max=x1, x2_min=x1)

        # k=2 (表示取第2个交点)
        x2, y2 = solve_two_parabolics_intersection_point(b2, b1, b0, a2, a1, a0, k=2)
        print("第2个交点:", x2, y2)
        # 绘制抛物线
        plot_parabolic(b2, b1, b0, a2, a1, a0, points=[(x2, y2)], x1_max=x2, x2_min=x2)

        # 绘制完整抛物线
        plot_parabolic(b2, b1, b0, a2, a1, a0, points=[(x1, y1), (x2, y2)], y_min=-100)

    except Exception as e:
        # 打印异常信息
        print(e)


# 主函数
# 注: 请根据需要, 打开/关闭相关测试函数的注释
def main():
    # 测试： 两条正反抛物线切的情况
    test1()

    # 测试: 两条正反抛物线不相交的情况
    # test2()

    # 测试: 两条正反抛物线相交的情况
    # test3()


# 程序入口
if __name__ == "__main__":
    main()
